import sys,matplotlib
from matplotlib import colors, ticker, cm
sys.path.append('../../utils/')
matplotlib.use('Agg')
from numpy import *
from pylab import *
import getopt
import os, re
import pdb
import math as math
import linecache
from matplotlib.colors import LogNorm
from scipy import interpolate
import pickle as pickle

############################################################################################ 

# SPECIFY FILES + SPECIES TO PLOT 

#files = ['kinmars_old.out-1']
#files = ['lat_alt_runs_reruns/kinmars_H2O_02.out-1','lat_alt_runs_reruns/kinmars_H2O_08.out-1','lat_alt_runs_reruns/kinmars_H2O_13.out-1']
#files = ['kinmars_old.out-1','kinmars_old_1.out-1','kinmars_changes.out-1']
files = ['lat_alt_runs_reruns/kinmars_H2O_02.out-1','lat_alt_runs_reruns/kinmars_H2O_54.out-1']
# ignore this ,'kinmars.out-1']

#labels = ['frank kindata','replaced H2O file']
#labels = ['arthur kindata','ranjan','frank kindata']
labels = ['latitude 2','latitude 54']
# ignore this 'our kindata']
lns = ['-','--','-.',':'] # linestyle (solid, dash, dash-dot, dotted). One per filename

#pnames = ['CL','CL2','CLO','HCL','HOCL','CLCO','COCL2','CLCO3',\
#          'OCLO', 'CLO3', 'CL2O', 'CL2O2', 'CL2O3', 'CLNO','CLNO2','CLONO','CLNO3']
#pnames = ['CO','O2','H2','O3','NO','HNO3','CH4','OH','HO2','H']
#pnames = ['NO','NO2','NO3','HNO2','HNO3','HO2NO2','N2O']
#pnames = ['OH','HO2','H2O2']
#pnames = ['CL','CL2','CLO','HCL','HOCL','CLCO3','CLNO','CLO3','CLNO3']
#pnames = ['O', 'O2', 'O3','CO', 'OH','H2','HNO3','NO','HO2']
pnames = ['CO','O2','H2','O3','NO','HNO3','OH','HO2','H']
 # can take more/less species, just make sure len(pnames) < len(pclr)
pclr = ['magenta','purple','navy','blue','orange','green','red','cyan','brown',\
        'magenta','purple','navy','blue','orange','green','red','cyan','brown'] # one color per species

############################################################################################

# Plot Style Things
matplotlib.rc('font', family='serif')
matplotlib.rc('xtick', labelsize = 25)
matplotlib.rc('ytick', labelsize = 25)
matplotlib.rc('axes', linewidth=2)
fontname = 'serif'
fontsize=30
lw = 3.3

# OPEN FIGURE 
fig,ax=plt.subplots(figsize=(12,12))

############################################################################################

# READ & PLOT FOR EACH FILE

# Get data from "files", loop over filenames in list
fctr = 0
for ifile in files:
    infile = open(ifile,'r')
    temp = infile.readline()

    while not 'NATOM' in temp:
        temp = infile.readline()
    temp = infile.readline().split() # vals 
    natom, nmol, nreact, nfix, nvary = int(temp[0]), int(temp[1]), int(temp[2]), int(temp[4]), int(temp[6])
    temp = infile.readline() # header       
    temp = infile.readline().split() # vals 
    nz = int(temp[2])

    # skip all lines that come before 'MIXING RATIO'
    temp = infile.readline()
    while 'MIXING RATIO' not in temp:
        temp = infile.readline()
    
    # loop over groups of species 
    mix = {}
    alt = np.zeros(nz)

    blank = infile.readline()
    header = infile.readline().split()
    alt = np.zeros(nz)
    while len(header)>1:
        for ix in range(1,len(header)):
            mix[header[ix]]= np.zeros(nz)
        for iz in range(0,nz):
            temp = infile.readline().split()
            for ix in range(1,len(header)):
                mix[header[ix]][iz] = float(temp[ix+1])
            alt[iz] = float(temp[1])
        blank = infile.readline()
        header = infile.readline().split()

    for ipnames in range(0,len(pnames)):
   #     if not mix.has_key(pnames[ipnames]): # ensures there are no errors if a species is not in the model
   #             continue
        if fctr == 0: # create "label" only for first file - this is the legend that will name colors with species
            ax.semilogx(mix[pnames[ipnames]],alt,color=pclr[ipnames],linestyle=lns[fctr],linewidth=lw,label=pnames[ipnames])
        else: # after first file, plot but don't add to the legend
            ax.semilogx(mix[pnames[ipnames]],alt,color=pclr[ipnames],linestyle=lns[fctr],linewidth=lw)

    # plot a dot in black way outside the limits just to get an extra label for this linestyle/filename
    ax.semilogx([1e-30,1e-30],[0,0],color='k',linestyle=lns[fctr],label=labels[fctr])
    fctr +=1 

# put legend at top of plot, use 5 columns to list everything
ax.legend(ncol=8)

# axis limits 
ax.set_ylim([0,200])
ax.set_xlim([1e-15,100e1])
ax.set_xlabel('Mixing Ratio',fontsize=fontsize,fontname=fontname)
ax.set_ylabel('Altitude',fontsize=fontsize,fontname=fontname)

plt.savefig('mixingOyCOyHy_1D_aug_24.png')

#bbox_to_anchor=(0.54,1.1)
